import argparse

import numpy as np
from PIL import Image
import torch,torch_npu
from torchvision import transforms
from models.get_model import get_arch
from utils.model_saving_loading import load_model
import os,json
# argument parsing
parser = argparse.ArgumentParser()
required_named = parser.add_argument_group('required arguments')
parser.add_argument('--data_path', help='path to dataset',
                    default='../Dataset/infantImages')
parser.add_argument('--model_name', help='model used for segmentation',
                    default='wnet')
parser.add_argument('--bin_thresh', type=float,
                    default='0.4196', help='binarizing threshold')
parser.add_argument('--result_path', type=str, default=None,
                    help='path to save prediction)')


class Vascular_segmentation:
    def __init__(self, model_name, resize=(452,600), bin_threshold='0.4196') -> None:
        if torch.npu.is_available():
            self.device = torch.device('npu')
        elif torch.cuda.is_available():
            self.device = torch.device('cuda')
        else:
            self.device = torch.device('cpu')
            
        self.bin_thresh = bin_threshold
        self.resize = resize
        self.image_process = transforms.Compose(
            [transforms.Resize(resize), transforms.ToTensor()])
        if model_name == 'wnet':
            model_save_path='./checkpoints/wnet_drive/'
        elif model_name == 'big_wnet_1024':
            model_save_path='./checkpoints/big_wnet_hrf_av_1024/'
            model_name='big_wnet'
            raise NotImplementedError
        elif model_name == 'big_wnet':
            model_save_path='./checkpoints/big_wnet_drive_av/'
            raise NotImplementedError
        
        model = get_arch(model_name).to(self.device)
        self.model, _ = load_model(model, model_save_path, self.device)
        self.model.eval()
    def _vasular_segmentation(self, image_path, mask_path=None, save_path=None, bin_save_path=None):
        # Load and apply mask if provided
        if mask_path is not None:
            mask = Image.open(mask_path).convert('L').resize((self.resize[1],self.resize[0]))
            mask = np.array(mask)
            mask[mask > 0] = 1
        
        # Process image
        img = Image.open(image_path)
        img = self.image_process(img)
        
        with torch.no_grad():
            img = img.unsqueeze(0).to(self.device)
            pred_img = self.model(img).cpu().squeeze()
            pred_img = torch.sigmoid(pred_img).numpy()
            
            # Apply mask to prediction if provided
            if mask_path is not None:
                pred_img = pred_img * mask
            
            # Save the prediction result as an image if save_path is provided
            if save_path is not None:
                pred_img_img = Image.fromarray((pred_img * 255).astype(np.uint8))
                pred_img_img.save(save_path)
            
            # Save binary thresholded prediction if bin_save_path is provided
            if bin_save_path is not None:
                binary_img = (pred_img > self.bin_thresh).astype(np.uint8) * 255
                binary_img = Image.fromarray(binary_img)
                binary_img.save(bin_save_path)
        
        return pred_img

if __name__ == '__main__':
    from tqdm import tqdm
    args = parser.parse_args()
    segmentation_handler = Vascular_segmentation(args.model_name, bin_threshold=args.bin_thresh)

    with open(os.path.join(args.data_path, 'annotations.json')) as f:
        data_dict = json.load(f)
        
    vascular_dir = os.path.join(args.data_path, 'vascular')
    os.makedirs(vascular_dir, exist_ok=True)

    for image_name in tqdm(data_dict, desc="Processing Images", unit="image"):
        data = data_dict[image_name]
        image_path = os.path.join(args.data_path, data['enhanced_path'])
        vascular_path = os.path.join(vascular_dir, image_name)
        
        if os.path.exists(vascular_path):
            if 'vascular_path' not in data:
                data_dict[image_name]['vascular_path'] = os.path.join('vascular', image_name)
            continue
        
        if 'mask_path' in data:
            mask_path = os.path.join(args.data_path, data['mask_path'])
            pred_img = segmentation_handler._vasular_segmentation(image_path, mask_path, save_path=vascular_path)
        else:
            pred_img = segmentation_handler._vasular_segmentation(image_path, save_path=vascular_path)
        data_dict[image_name]['vascular_path'] = os.path.join('vascular', image_name)

    with open(os.path.join(args.data_path,'annotations.json'),'w') as f:
        json.dump(data_dict,f)